# -*- coding: utf-8 -*-
# created on 2016/11/30
import copy

from mathsolver.functions.shulie.ShuLie_getSeqPoly import ShuLieGetSeqPoly005, ShuLieGetSeqPoly002
from mathsolver.functions.shulie.ShuLie_Transform import (DengChaBiEqNormalization, DengChaBiExprNormalization,
                                                          DengChaBiJiChuAssume)
from mathsolver.functions.shulie.basic import *
from mathsolver.functions.shulie.base import *
from mathsolver.functions.mathematica.mathematicaSolve import MathematicaSolve
from mathsolver.functions.base import *
from sympy import sympify
from sympy.abc import n


class ShuLieGetTongXiang001(BaseFunction):
    """
    1.求数列的通项
    2.单数列
    3.等差，等比数列型
    4.new_otherSeqs空
    eg. 已知数列{a_{n}}为等差数列,a_{1}=2,a_{1}+a_{2}+a_{3}=12.求数列{a_{n}}的通项公式
    """
    def solver(self, *args):
        assert len(args) == 1
        known = dict(self.known)
        assert 'Sequations' in known
        assert known['Sequations']
        seqs = known['Sequations']
        assert isinstance(seqs, list)
        if known['inequations']:
            for ineq in known['inequations']:
                left, op, right = ineq
                new_ineq = [sympify(left), op, sympify(right)]
                if new_ineq not in seqs:
                    seqs.append(new_ineq)
        if known['equations']:
            for eq in known['equations']:
                if eq not in seqs:
                    seqs.append(eq)
        new_seqs, new_other_seqs = sl_seqs_other_seqs(known, seqs)
        assert not new_other_seqs
        assert new_seqs
        sls, sl_flag = sl_shulies_flag(known, new_seqs)
        assert sum(sl_flag) == 0
        stepsolver = DengChaBiJiChuAssume(known=self.known).solver(sls)
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        stepsolver = DengChaBiEqNormalization(known=self.known).solver(seqs)
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        stepsolver = MathematicaSolve().solver(stepsolver.output[0], [])
        self.steps.append(["", "解得: %s" % (stepsolver.output[0].printing())])
        solutions = stepsolver.output[0]
        target = args[0].sympify()
        stepsolver = DengChaBiExprNormalization(known).solver(BasePoly(target))
        self.label.update(stepsolver.label)
        self.steps.append(["", "%s = %s" % (new_latex(target), new_latex(stepsolver.output[0].sympify()))])
        new_target = stepsolver.output[0].sympify()
        if solutions.sympify():
            if type(solutions.sympify()) == bool:
                target_values = [new_target]
            else:
                target_values = sl_solutuons_subs(solutions, new_target)
        else:
            target_values = [new_target]
        self.steps.append(["", "∴%s" % "或".join(["%s = %s" % (new_latex(args[0].sympify()), new_latex(v)) for v in target_values])])
        answer = {target: target_values}
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieGetTongXiang003(BaseFunction):
    """
    已知{a_{n}}是公差为3的等差数列,数列{b_{n}}满足b_{1}=1,b_{2}=\\frac{1}{3},a_{n}b_{n+1}+b_{n+1}=nb_{n}.求{a_{n}}的通项公式.
    1.求数列的通项
    2.双数列: 一个数列为等差，等比类型; 另外一个为一般数列
    3.所求数列为等差，等比数列
    """
    def solver(self, *args):
        assert len(args) == 1
        known = self.known
        assert 'Sequations' in known
        assert known['Sequations']
        seqs = known['Sequations']
        new_seqs, new_other_seqs = sl_seqs_other_seqs(self.known, seqs)
        assert not new_other_seqs
        assert new_seqs
        seq = args[0].sympify()
        sls, sl_flag = sl_shulies_flag(self.known, new_seqs)
        assert sum(sl_flag) > 0
        assert len(sls) == 2
        seq_name = seq.func
        seq_sl = self.search(str(seq_name))
        seq_sl_type = seq_sl.type
        assert (seq_sl_type == str("DengCha") or seq_sl_type == str("DengBi"))
        ditui_eqs, nditui_eqs = sl_ditui_or_not_ditui(seqs)
        condition_eqs = []
        condition_eqs_symbols = []
        for NDiTui_Eq in nditui_eqs:
            if len(NDiTui_Eq) == 2:
                expr = NDiTui_Eq[0] - NDiTui_Eq[1]
            else:
                expr = NDiTui_Eq[0] - NDiTui_Eq[2]
            expr = ShuLieSnExprNormalization().solver(BasePoly(expr), sympify('a_')).output[0].sympify()
            expr_xiangs = sl_xiangs(expr)
            for expr_xiang in expr_xiangs:
                condition_eqs_symbols.append(expr_xiang)
            condition_eqs.append(NDiTui_Eq)
        condition_eqs_num = len(condition_eqs)
        condition_eqs_symbols = set(condition_eqs_symbols)
        condition_eqs_symbols = list(condition_eqs_symbols)
        for eq in ditui_eqs:
            if len(eq) == 2:
                expr = eq[0] - eq[1]
            else:
                expr = eq[0] - eq[2]
            for i in range(1, 5):
                new_expr = expr.subs({seq_sl.nName: i})
                new_expr = ShuLieSnExprNormalization().solver(BasePoly(new_expr), sympify('a_')).output[0].sympify()
                if len(condition_eqs_symbols) == condition_eqs_num and condition_eqs_num != 0:
                    break
                else:
                    if len(eq) == 2:
                        condition_eqs.append([new_expr, S.Zero])
                    else:
                        condition_eqs.append([new_expr, eq[1], S.Zero])
                    condition_eqs_num += 1
                    expr_xiangs = sl_xiangs(new_expr)
                    for expr_xiang in expr_xiangs:
                        condition_eqs_symbols.append(expr_xiang)
                    condition_eqs_symbols = set(condition_eqs_symbols)
                    condition_eqs_symbols = list(condition_eqs_symbols)
        stepsolver1 = ShuLieDengChaBiSolveEqs(known).solver(BaseIneqs(condition_eqs), [])
        self.steps += stepsolver1.steps
        self.label.update(stepsolver1.label)
        solutions = stepsolver1.output[0]
        stepsolver2 = DengChaBiExprNormalization(known).solver(BasePoly(seq))
        self.steps += stepsolver2.steps
        new_seq = stepsolver2.output[0].sympify()
        expr_xiangs = sl_xiangs(new_seq)
        for expr_xiang in expr_xiangs:
            condition_eqs_symbols.append(expr_xiang)
        for eq in ditui_eqs:
            if len(eq) == 2:
                expr = eq[0] - eq[1]
            else:
                expr = eq[0] - eq[2]
            for i in range(1, 5):
                new_expr = expr.subs({seq_sl.nName: i})
                new_expr = ShuLieSnExprNormalization().solver(BasePoly(new_expr), sympify('a_')).output[0].sympify()
                if len(condition_eqs_symbols) == condition_eqs_num and condition_eqs_num != 0:
                    break
                else:
                    if len(eq) == 2:
                        condition_eqs.append([new_expr, S.Zero])
                    else:
                        condition_eqs.append([new_expr, eq[1], S.Zero])
                    condition_eqs_num += 1
                    expr_xiangs = sl_xiangs(new_expr)
                    for expr_xiang in expr_xiangs:
                        condition_eqs_symbols.append(expr_xiang)
                    condition_eqs_symbols = set(condition_eqs_symbols)
                    condition_eqs_symbols = list(condition_eqs_symbols)
        stepsolver = MathematicaSolve().solver(BaseEqs(condition_eqs))
        solutions = stepsolver.output[0]
        seq_values = sl_solutuons_subs(solutions, new_seq)
        self.steps.append(["", "∴%s" % "或".join(["%s = %s" % (new_latex(seq), new_latex(v)) for v in seq_values])])
        answer = {seq: seq_values}
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieGetTongXiang004(BaseFunction):
    """
    设等差数列{a_{n}的公差为d,前n项和为S_{n},等比数列{b_{n}}的公比为q.已知b_{1}=a_{1},b_{2}=2,q=d,S_{10}=100.当d>1时,记c_{n}=\\frac{a_{n}}{b_{n}},求数列{c_{n}}的前n项和为T_{n}.
    1.求数列的项或和
    2.单数列
    3.等差，等比数列型
    4.new_otherSeqs非空
    """
    def solver(self, *args):
        assert len(args) == 1
        known = self.known
        assert 'Sequations' in known
        assert known['Sequations']
        seqs = known['Sequations']
        new_seqs, new_other_seqs = sl_seqs_other_seqs(self.known, seqs)
        assert new_other_seqs
        assert new_seqs
        seq = args[0].sympify()
        seq_values = sl_solve_an(new_other_seqs, seq)
        assert len(seq_values) == 1
        new_seq = seq_values[0]
        known['Sequations'] = new_seqs
        stepsolver = ShuLieGetSeqPoly002(known).solver(BaseSeqPoly(new_seq))
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        new_seq_values = stepsolver.output[0].sympify()
        answer = {seq: new_seq_values}
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieGetTongXiang005(BaseFunction):
    """
    已知数列{a_{n}}中,a_{1}=2,{a}_{n+1}=(\\sqrt{2}-1)({a}_{n}+2),求{a_{n}}的通项公式
    1.求数列的通项
    2.单数列
    3.一般数列型
    4.new_otherSeqs空
    """
    def solver(self, *args):
        assert len(args) == 1
        known = self.known
        assert 'Sequations' in known
        assert known['Sequations']
        seqs = known['Sequations']
        if self.known['inequations']:
            for ineq in self.known['inequations']:
                left, op, right = ineq
                new_ineq = [sympify(left), op, sympify(right)]
                if new_ineq not in seqs:
                    seqs.append(new_ineq)
        if self.known['equations']:
            for eq in self.known['equations']:
                if eq not in seqs:
                    seqs.append(eq)
        new_seqs, new_other_seqs = sl_seqs_other_seqs(self.known, seqs)
        assert not new_other_seqs
        assert new_seqs
        seq = args[0].sympify()
        sls, sl_flag = sl_shulies_flag(self.known, new_seqs)
        assert sum(sl_flag) > 0
        assert len(sls) == 1
        sl = sls[0]
        sl.eqs_update(new_seqs)
        sl.get_shouxiang()
        sl.get_dituigongshi()
        sl.get_xiangs()
        sl.get_nineqs()
        sl.get_tongxianggongshi()
        self.steps += sl.steps_TongXiang
        self.label.update(sl.label_TongXiang)
        seq_values = sl.an
        answer = {seq: seq_values}
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieGetTongXiang006(BaseFunction):
    """
    已知数列{a_{n}}中,a_{1}=1.a_{n+1}=\\frac{5}{2}-\\frac{1}{{a}_{n}},b_{n}=\\frac{1}{{a}_{n}-2},求数列{b_{n}}的通项公式.
    1.求数列的通项
    2.单数列
    3.一般数列型
    4.new_otherSeqs非空
    """
    def solver(self, *args):
        assert len(args) == 1
        known = self.known
        assert 'Sequations' in known
        assert known['Sequations']
        seqs = known['Sequations']
        new_seqs, new_other_seqs = sl_seqs_other_seqs(self.known, seqs)
        assert new_other_seqs
        assert new_seqs
        sls, sl_flag = sl_shulies_flag(self.known, new_seqs)
        assert sum(sl_flag) > 0
        assert len(sls) == 1
        seq = args[0].sympify()
        seq_values = sl_solve_an(new_other_seqs, seq)
        assert len(seq_values) == 1
        new_seq = seq_values[0]
        known['Sequations'] = new_seqs
        stepsolver = ShuLieGetSeqPoly005(known).solver(BaseSeqPoly(new_seq))
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        new_seq_values = stepsolver.output[0].sympify()
        answer = {seq: new_seq_values}
        self.output.append(BaseShuLieValue(answer))
        self.steps.append(["", "∴ %s" % ",".join(["%s = %s" % (new_latex(seq), new_latex(v)) for v in new_seq_values])])
        return self


class ShuLieGetTongXiang007(BaseFunction):
    """
    # 已知{a_{n}}是公差为3的等差数列,数列{b_{n}}满足b_{1}=1,b_{2}=\\frac{1}{3},a_{n}b_{n+1}+b_{n+1}=nb_{n}.求{b_{n}}的通项公式;
    # 1.求数列的通项
    # 2.双数列: 一个数列为等差，等比类型; 另外一个为一般数列
    # 3.所求数列为一般数列
    """
    def solver(self, *args):
        assert len(args) == 1
        known = self.known
        assert 'Sequations' in known
        seqs = known['Sequations']
        assert seqs
        new_seqs, new_other_seqs = sl_seqs_other_seqs(known, seqs)
        assert not new_other_seqs
        assert new_seqs
        sls, sl_flag = sl_shulies_flag(known, new_seqs)
        assert sum(sl_flag) > 0
        assert len(sls) == 2
        seq = args[0].sympify()
        otherseqs = []
        for eq in new_seqs:
            xiangs = sl_xiangs001(eq[0] - eq[1])
            if xiangs:
                for xiang in xiangs:
                    xiang = sympify(xiang)
                    xiang_name = xiang.func
                    if str(xiang_name) != str(seq.func):
                        otherseqs.append(xiang_name(sympify('n')))
        otherseqs = set(otherseqs)
        otherseqs = list(otherseqs)
        assert len(otherseqs) == 1
        otherseq = otherseqs[0]
        otherseq_name = otherseq.func
        otherseq_sl = self.search(str(otherseq_name))
        otherseq_sl_type = otherseq_sl.type
        assert (otherseq_sl_type == "DengCha" or otherseq_sl_type == "DengBi")
        stepsolver = DengChaBiJiChuAssume(known=self.known).solver([otherseq_sl])
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        stepsolver = ShuLieGetTongXiang003(known=self.known).solver(BasePoly(otherseq))
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        tongxiangs = stepsolver.output[0].sympify()
        tongxiangs = tongxiangs[otherseq]
        assert len(tongxiangs) == 1
        tongxiang = tongxiangs[0]
        new_known = dict(known)
        update_seqs = []
        for eq in new_seqs:
            expr = eq[0] - eq[1]
            expr = simplify(expr)
            xiangs = sl_xiangs001(expr)
            if xiangs:
                for xiang in xiangs:
                    xiang = sympify(xiang)
                    xiang_name = xiang.func
                    if str(xiang_name) == str(otherseq.func):
                        xiang_index = xiang.args[0]
                        xiang_value = tongxiang.subs({sympify('n'): xiang_index})
                        expr = expr.subs({xiang: xiang_value})
                        expr = simplify(expr)
            update_seqs.append([expr, S.Zero])
        new_known['Sequations'] = update_seqs
        stepsolver = ShuLieGetTongXiang005(new_known).solver(BasePoly(seq))
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        seq_values = stepsolver.output[0].sympify()
        answer = {seq: seq_values[seq]}
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieGetTongXiang008(BaseFunction):
    """
    已知等比数列{a_{n}}中,a_{1}=\\frac{1}{3},公比q=\\frac{1}{3}.设b_{n}=log_{3}a_{1}+log_{3}a_{2}+…+log_{3}a_{n},求数列{b_{n}}的通项公式.
    1. 求数列的通项
    2. 单数列
    4. b_{n}=log_{3}a_{1}+log_{3}a_{2}+…+log_{3}a_{n}
    """
    def solver(self, *args):
        assert len(args) == 1
        known = self.known
        assert 'Sequations' in known
        assert known['Sequations']
        Seqs = known['Sequations']
        new_Seqs, new_otherSeqs = sl_Seqs_OtherSeqs(known, Seqs)
        assert not new_otherSeqs
        assert new_Seqs
        assert 'InfSequations' in known
        assert known['InfSequations']
        InfSeqs = known['InfSequations']
        assert len(InfSeqs) == 1
        InfSeq = InfSeqs[0]
        left, right = InfSeq
        new_known = dict(known)
        del new_known['InfSequations']
        if isinstance(left, str):
            raise Exception("to do")
        else:
            from mathsolver.functions.shulie.ShuLie_getInfSums import ShuLie_InfSums
            stepsolver = ShuLie_InfSums(new_known).solver(baseSeqinfpoly(right))
            self.steps += stepsolver.steps
            self.label.update(stepsolver.label)
            seq_values = stepsolver.output[0].sympify()
            answer = {}
            answer[left] = seq_values
            self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieGetTongXiang009(BaseFunction):
    """
    在数列{a_{n}}中,a_{1}=1,a_{n+1}=(1+\\frac{1}{n})a_{n}+\\frac{n+1}{{2}^{n}}.设b_{n}=\\frac{{a}_{n}}{n},求数列{b_{n}}的通项公式;
    """

    def solver(self, *args):
        assert len(args) == 1
        known = self.known
        assert 'Sequations' in known
        assert known['Sequations']
        Seqs = known['Sequations']
        new_Seqs, new_otherSeqs = sl_Seqs_OtherSeqs(self.known, Seqs)
        assert new_otherSeqs
        assert new_Seqs
        ShuLies, ShuLie_flag = sl_ShuLies_ShuLiesflag(self.known, new_Seqs)
        assert sum(ShuLie_flag) > 0
        assert len(ShuLies) == 1
        ShuLie = ShuLies[0]
        ShuLie.get_Init()
        ShuLie.Eqs_Update(new_Seqs)
        ShuLie.Eqs_Update(new_otherSeqs)
        ShuLie.get_ShouXiang()
        ShuLie.get_DiTuiGongShi()
        ShuLie.get_Xiangs()
        ShuLie.get_nXiangs()
        ShuLie.get_nIneqs()
        aName = ShuLie.aName
        nName = ShuLie.nName
        seq = args[0].sympify()
        seq_values = []
        if str(seq.func) == str(aName):
            if not ShuLie.an:
                ShuLie.get_TongXiangGongShi()
                self.steps += ShuLie.steps_TongXiang
                self.label.update(ShuLie.label_TongXiang)
            seq_values = ShuLie.an
        else:
            new_seqs = sl_solve_an(new_otherSeqs, seq)
            assert len(new_seqs) == 1
            new_seq = new_seqs[0]
            new_seq_names = sl_ShuLie_Name(new_seq)
            new_seq_xiangs = sl_Xiangs(new_seq)
            if len(new_seq_names) == 1:
                new_seq_name = new_seq_names[0]
                if str(new_seq_name) == str(aName):
                    if not ShuLie.an:
                        ShuLie.get_TongXiangGongShi()
                        self.steps += ShuLie.steps_TongXiang
                        self.label.update(ShuLie.label_TongXiang)
                    TongXiangs = ShuLie.an
                    for TongXiang in TongXiangs:
                        seq_value = new_seq
                        for new_seq_xiang in new_seq_xiangs:
                            new_seq_xiang = sympify(new_seq_xiang)
                            new_seq_xiang_index = new_seq_xiang.args[0]
                            new_seq_xiang_value = TongXiang.subs({n: new_seq_xiang_index})
                            seq_value = seq_value.subs({new_seq_xiang: new_seq_xiang_value})
                        seq_value = seq_value.expand().simplify()
                        seq_values.append(seq_value)
        seq_values = set(seq_values)
        seq_values = list(seq_values)
        assert seq_values
        self.steps.append(["", "∴ %s" % (",".join(["%s = %s" % (new_latex(seq), new_latex(value)) for value in seq_values]))])
        answer = {}
        answer[seq] = seq_values
        self.output.append(BaseShuLieValue(answer))
        return self


class ShuLieGetTongXiang(BaseFunction):
    CLS = [ShuLieGetTongXiang001, ShuLieGetTongXiang003, ShuLieGetTongXiang004,
           ShuLieGetTongXiang005, ShuLieGetTongXiang006, ShuLieGetTongXiang007, ShuLieGetTongXiang008, 
           ShuLieGetTongXiang009]

    def solver(self, *args):
        known = self.known
        r = None
        for cl in ShuLieGetTongXiang.CLS:
            try:
                new_known = copy.deepcopy(known)
                r = cl(new_known, verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            print('ShuLieGetTongXiang failed')
        return r


class ShuLieGetTongXiangs(BaseFunction):
    def solver(self, *args):
        seqs = args[0].sympify()
        answer = {}
        for seq in seqs:
            new_known = dict(self.known)
            stepsolver = ShuLieGetTongXiang(new_known).solver(BaseSequence(seq))
            self.steps += stepsolver.steps
            tong_xiangs = stepsolver.output[0].sympify()
            answer[seq] = tong_xiangs
        self.output.append(BaseShuLieValue(answer))
        return self
